import sys
import warnings
from PyQt5.QtWidgets import (
    QApplication, QMainWindow, QVBoxLayout, QHBoxLayout,
    QPushButton, QWidget, QOpenGLWidget, QSlider, QLabel,
    QStackedWidget
)
from PyQt5.QtCore import Qt, QTimer
from PyQt5.QtGui import QSurfaceFormat
from trajectory_animation import TrajectoryAnimation
from OpenGL.GL import *

warnings.filterwarnings("ignore", category=DeprecationWarning)

fmt = QSurfaceFormat()
fmt.setSamples(8)
QSurfaceFormat.setDefaultFormat(fmt)

class OpenGLWidget(QOpenGLWidget):
    def __init__(self, animation, parent=None):
        super().__init__(parent)
        self.animation = animation
        self.last_mouse_pos = None

    def initializeGL(self):
        glClearColor(0.0, 0.0, 0.0, 1.0)
        glEnable(GL_DEPTH_TEST)
        glEnable(GL_LINE_SMOOTH)
        glHint(GL_LINE_SMOOTH_HINT, GL_NICEST)
        glEnable(GL_POINT_SMOOTH)
        glHint(GL_POINT_SMOOTH_HINT, GL_NICEST)
        glHint(GL_POLYGON_SMOOTH_HINT, GL_NICEST)
        self.animation.setup_view(self.width(), self.height())

    def resizeGL(self, w, h):
        self.animation.setup_view(w, h)

    def paintGL(self):
        self.animation.draw(self.width(), self.height())

    def keyPressEvent(self, event):
        key = event.text().encode('utf-8')
        self.animation.keyboard(key)
        self.update()

    def mousePressEvent(self, event):
        if event.button() == Qt.LeftButton:
            self.last_mouse_pos = event.pos()

    def mouseMoveEvent(self, event):
        if self.last_mouse_pos and event.buttons() & Qt.LeftButton:
            dx = event.x() - self.last_mouse_pos.x()
            dy = event.y() - self.last_mouse_pos.y()
            self.animation.angle_y += dx * 0.5
            self.animation.angle_x += dy * 0.5
            self.last_mouse_pos = event.pos()
            self.update()

    def mouseReleaseEvent(self, event):
        self.last_mouse_pos = None

class MainWindow(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("PyQt + OpenGL 3D轨迹动画")
        self.setGeometry(50, 50, 1600, 1000)
        self.setStyleSheet("background-color: #001f3f; color: white;")

        points_file = "points.csv"
        self.animation = TrajectoryAnimation(points_file)
        self.opengl_widget_main = OpenGLWidget(self.animation)
        self.opengl_widget_position = OpenGLWidget(self.animation)
        self.animation_active = False

        # 页面 0：主界面
        self.page_main = QWidget()
        main_layout = QVBoxLayout()
        content_layout = QHBoxLayout()
        content_layout.addWidget(self.opengl_widget_main)
        main_layout.addLayout(content_layout, stretch=8)

        bottom_button_layout = QHBoxLayout()
        btn_position = QPushButton("位置")
        btn_teach = QPushButton("示教")
        btn_replay = QPushButton("复现")
        btn_execute = QPushButton("执行")
        for btn, color in zip([btn_position, btn_teach, btn_replay, btn_execute],
                               ["#333333", "#0000FF", "#FFA500", "#008000"]):
            btn.setStyleSheet(f"background-color: {color}; color: white;")
        btn_replay.clicked.connect(self.start_animation)
        btn_execute.clicked.connect(self.start_animation)
        btn_position.clicked.connect(self.show_position_page)

        bottom_button_layout.addWidget(btn_position)
        bottom_button_layout.addStretch(1)
        bottom_button_layout.addWidget(btn_teach)
        bottom_button_layout.addWidget(btn_replay)
        bottom_button_layout.addWidget(btn_execute)
        bottom_button_layout.addStretch(1)

        self.process_labels = []
        process_steps = ["毛坯准备", "划线", "钻孔", "铰孔", "攻丝",
                         "铣削", "车削", "磨削", "热处理", "装配", "检验"]
        process_layout = QHBoxLayout()
        for i, step in enumerate(process_steps):
            label = QLabel(step)
            label.setStyleSheet(f"""
                background-color: {'#cc0000' if i == 0 else '#004080'};
                color: white;
                font-family: SimSun;
                padding: 6px 12px;
                border: 2px solid {'#ff3333' if i == 0 else '#3399ff'};
                border-radius: 4px;
                qproperty-alignment: 'AlignCenter';
            """)
            process_layout.addWidget(label)
            self.process_labels.append(label)
        process_container = QWidget()
        process_container.setLayout(process_layout)

        slider_layout = QVBoxLayout()
        self.scale_slider = QSlider(Qt.Horizontal)
        self.scale_slider.setMinimum(2)
        self.scale_slider.setMaximum(50)
        self.scale_slider.setValue(10)
        self.scale_slider.setTickInterval(1)
        self.scale_slider.setTickPosition(QSlider.TicksBelow)
        self.scale_slider.valueChanged.connect(self.update_scale)
        slider_label = QLabel("缩放")
        slider_label.setStyleSheet("color: white;")
        slider_layout.addWidget(slider_label)
        slider_layout.addWidget(self.scale_slider)

        main_layout.addLayout(bottom_button_layout, stretch=1)
        main_layout.addWidget(process_container)
        main_layout.addLayout(slider_layout, stretch=1)
        self.page_main.setLayout(main_layout)

        # 页面 1：坐标页
        self.page_position = QWidget()
        pos_layout = QVBoxLayout()

        coord_view_layout = QHBoxLayout()
        coord_view_layout.addWidget(self.opengl_widget_position)
        pos_layout.addLayout(coord_view_layout, stretch=8)

        self.label_x = QLabel("X: ---")
        self.label_y = QLabel("Y: ---")
        self.label_z = QLabel("Z: ---")

        for lbl, color in zip([self.label_x, self.label_y, self.label_z],
                               ["#00FF00", "#FFFFFF", "#3399FF"]):
            lbl.setStyleSheet(f"""
                background-color: {color};
                color: black;
                font-size: 24px;
                padding: 12px;
                border: 2px solid black;
                border-radius: 6px;
            """)
            lbl.setAlignment(Qt.AlignCenter)
            lbl.setFixedWidth(300)
            lbl.setFixedHeight(60)

        coord_layout = QHBoxLayout()
        coord_layout.setSpacing(20)
        coord_layout.addStretch(1)
        coord_layout.addWidget(self.label_x)
        coord_layout.addWidget(self.label_y)
        coord_layout.addWidget(self.label_z)
        coord_layout.addStretch(1)

        btn_back = QPushButton("返回")
        btn_back.setStyleSheet("background-color: #555555; color: white;")
        btn_back.clicked.connect(self.show_main_page)
        btn_back.setFixedWidth(120)
        btn_back.setFixedHeight(40)

        pos_layout.addLayout(coord_layout)
        pos_layout.addWidget(btn_back, alignment=Qt.AlignCenter)

        self.page_position.setLayout(pos_layout)

        # 页面容器
        self.stack = QStackedWidget()
        self.stack.addWidget(self.page_main)
        self.stack.addWidget(self.page_position)
        self.setCentralWidget(self.stack)

        # 动画更新定时器
        self.prev_process_step = 0
        self.timer = QTimer()
        self.timer.timeout.connect(self.update_animation)
        self.timer.start(self.animation.timer_interval)

    def show_position_page(self):
        self.stack.setCurrentIndex(1)

    def show_main_page(self):
        self.stack.setCurrentIndex(0)

    def start_animation(self):
        self.animation_active = True

    def update_animation(self):
        if self.animation_active:
            self.animation.update_animation()
            self.opengl_widget_main.update()
            self.opengl_widget_position.update()

            # 实时刷新坐标显示页面的坐标值
            if self.stack.currentIndex() == 1:
                index = self.animation.current_index
                if 0 <= index < len(self.animation.points):
                    x, y, z = self.animation.points[index]
                    self.label_x.setText(f"X: {x:.3f}")
                    self.label_y.setText(f"Y: {y:.3f}")
                    self.label_z.setText(f"Z: {z:.3f}")

            last_index = len(self.animation.points) - 1
            if self.animation.current_index % 20 == 0 or self.animation.current_index == last_index:
                curr_step = (len(self.process_labels) - 1
                             if self.animation.current_index == last_index
                             else self.animation.current_index // 20)
                if curr_step < len(self.process_labels):
                    self.process_labels[self.prev_process_step].setStyleSheet("""
                        background-color: #004080;
                        color: white;
                        font-family: SimSun;
                        padding: 6px 12px;
                        border: 2px solid #3399ff;
                        border-radius: 4px;
                        qproperty-alignment: 'AlignCenter';
                    """)
                    self.process_labels[curr_step].setStyleSheet("""
                        background-color: #cc0000;
                        color: white;
                        font-family: SimSun;
                        padding: 6px 12px;
                        border: 2px solid #ff3333;
                        border-radius: 4px;
                        qproperty-alignment: 'AlignCenter';
                    """)
                    self.prev_process_step = curr_step

    def update_scale(self, value):
        self.animation.scale_factor = value / 10.0
        self.opengl_widget_main.update()
        self.opengl_widget_position.update()

if __name__ == "__main__":
    app = QApplication(sys.argv)
    window = MainWindow()
    window.show()
    sys.exit(app.exec_())
